// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <algorithm>
#include <cstdint>
#include <cstring>
#include <memory>
#include <sstream>
#include <string>

#include <ktest/ktest.h>

#include <nebula/array/array_nested.h>
#include <nebula/array/util.h>
#include <nebula/core/extension_type.h>
#include <nebula/io/memory.h>
#include <nebula/ipc/options.h>
#include <nebula/ipc/reader.h>
#include <nebula/ipc/writer.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/testing/extension_type.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/types/type.h>
#include <nebula/util/key_value_metadata.h>
#include <turbo/log/logging.h>

namespace nebula {
    class Parametric1Array : public ExtensionArray {
    public:
        using ExtensionArray::ExtensionArray;
    };

    class Parametric2Array : public ExtensionArray {
    public:
        using ExtensionArray::ExtensionArray;
    };

    // A parametric type where the extension_name() is always the same
    class Parametric1Type : public ExtensionType {
    public:
        explicit Parametric1Type(int32_t parameter)
            : ExtensionType(int32()), parameter_(parameter) {
        }

        int32_t parameter() const { return parameter_; }

        std::string extension_name() const override { return "parametric-type-1"; }

        bool ExtensionEquals(const ExtensionType &other) const override {
            const auto &other_ext = static_cast<const ExtensionType &>(other);
            if (other_ext.extension_name() != this->extension_name()) {
                return false;
            }
            return this->parameter() == static_cast<const Parametric1Type &>(other).parameter();
        }

        std::shared_ptr<Array> make_array(std::shared_ptr<ArrayData> data) const override {
            return std::make_shared<Parametric1Array>(data);
        }

        turbo::Result<std::shared_ptr<DataType> > deserialize(
            std::shared_ptr<DataType> storage_type,
            const std::string &serialized) const override {
            DKCHECK_EQ(4, serialized.size());
            const int32_t parameter = *reinterpret_cast<const int32_t *>(serialized.data());
            DKCHECK(storage_type->equals(int32()));
            return std::make_shared<Parametric1Type>(parameter);
        }

        std::string serialize() const override {
            std::string result("    ");
            memcpy(&result[0], &parameter_, sizeof(int32_t));
            return result;
        }

    private:
        int32_t parameter_;
    };

    // A parametric type where the extension_name() is different for each
    // parameter, and must be separately registered
    class Parametric2Type : public ExtensionType {
    public:
        explicit Parametric2Type(int32_t parameter)
            : ExtensionType(int32()), parameter_(parameter) {
        }

        int32_t parameter() const { return parameter_; }

        std::string extension_name() const override {
            std::stringstream ss;
            ss << "parametric-type-2<param=" << parameter_ << ">";
            return ss.str();
        }

        bool ExtensionEquals(const ExtensionType &other) const override {
            const auto &other_ext = static_cast<const ExtensionType &>(other);
            if (other_ext.extension_name() != this->extension_name()) {
                return false;
            }
            return this->parameter() == static_cast<const Parametric2Type &>(other).parameter();
        }

        std::shared_ptr<Array> make_array(std::shared_ptr<ArrayData> data) const override {
            return std::make_shared<Parametric2Array>(data);
        }

        turbo::Result<std::shared_ptr<DataType> > deserialize(
            std::shared_ptr<DataType> storage_type,
            const std::string &serialized) const override {
            DKCHECK_EQ(4, serialized.size());
            const int32_t parameter = *reinterpret_cast<const int32_t *>(serialized.data());
            DKCHECK(storage_type->equals(int32()));
            return std::make_shared<Parametric2Type>(parameter);
        }

        std::string serialize() const override {
            std::string result("    ");
            memcpy(&result[0], &parameter_, sizeof(int32_t));
            return result;
        }

    private:
        int32_t parameter_;
    };

    // An extension type with a non-primitive storage type
    class ExtStructArray : public ExtensionArray {
    public:
        using ExtensionArray::ExtensionArray;
    };

    class ExtStructType : public ExtensionType {
    public:
        ExtStructType()
            : ExtensionType(
                STRUCT({::nebula::field("a", int64()), ::nebula::field("b", float64())})) {
        }

        std::string extension_name() const override { return "ext-struct-type"; }

        bool ExtensionEquals(const ExtensionType &other) const override {
            const auto &other_ext = static_cast<const ExtensionType &>(other);
            if (other_ext.extension_name() != this->extension_name()) {
                return false;
            }
            return true;
        }

        std::shared_ptr<Array> make_array(std::shared_ptr<ArrayData> data) const override {
            return std::make_shared<ExtStructArray>(data);
        }

        turbo::Result<std::shared_ptr<DataType> > deserialize(
            std::shared_ptr<DataType> storage_type,
            const std::string &serialized) const override {
            if (serialized != "ext-struct-type-unique-code") {
                return turbo::invalid_argument_error("Type identifier did not match");
            }
            return std::make_shared<ExtStructType>();
        }

        std::string serialize() const override { return "ext-struct-type-unique-code"; }
    };

    class TestExtensionType : public ::testing::Test {
    public:
        void SetUp() { ASSERT_OK(RegisterExtensionType(std::make_shared<UuidType>())); }

        void TearDown() {
            if (GetExtensionType("uuid")) {
                ASSERT_OK(UnregisterExtensionType("uuid"));
            }
        }
    };

    TEST_F(TestExtensionType, ExtensionTypeTest) {
        auto type_not_exist = GetExtensionType("uuid-unknown");
        ASSERT_EQ(type_not_exist, nullptr);

        auto registered_type = GetExtensionType("uuid");
        ASSERT_NE(registered_type, nullptr);

        auto type = uuid();
        ASSERT_EQ(type->id(), Type::EXTENSION);
        ASSERT_EQ(type->bit_width(), 128);
        ASSERT_EQ(type->byte_width(), 16);

        const auto &ext_type = static_cast<const ExtensionType &>(*type);
        std::string serialized = ext_type.serialize();

        ASSERT_OK_AND_ASSIGN(auto deserialized,
                             ext_type.deserialize(fixed_size_binary(16), serialized));
        ASSERT_TRUE(deserialized->equals(*type));
        ASSERT_FALSE(deserialized->equals(*fixed_size_binary(16)));
        ASSERT_EQ(deserialized->id(), Type::EXTENSION);
        ASSERT_EQ(deserialized->bit_width(), 128);
        ASSERT_EQ(deserialized->byte_width(), 16);
    }

    auto RoundtripBatch = [](const std::shared_ptr<RecordBatch> &batch,
                             std::shared_ptr<RecordBatch> *out) {
        ASSERT_OK_AND_ASSIGN(auto out_stream, io::BufferOutputStream::create());
        ASSERT_OK(ipc::write_record_batch_stream({batch}, ipc::IpcWriteOptions::defaults(),
            out_stream.get()));

        ASSERT_OK_AND_ASSIGN(auto complete_ipc_stream, out_stream->finish());

        io::BufferReader reader(complete_ipc_stream);
        std::shared_ptr<RecordBatchReader> batch_reader;
        ASSERT_OK_AND_ASSIGN(batch_reader, ipc::RecordBatchStreamReader::open(&reader));
        ASSERT_OK(batch_reader->read_next(out));
    };

    TEST_F(TestExtensionType, IpcRoundtrip) {
        auto ext_arr = ExampleUuid();
        auto batch = RecordBatch::create(schema({field("f0", uuid())}), 4, {ext_arr});

        std::shared_ptr<RecordBatch> read_batch;
        RoundtripBatch(batch, &read_batch);
        CompareBatch(*batch, *read_batch, false /* compare_metadata */);

        // Wrap type in a ListArray and ensure it also makes it
        auto offsets_arr = assert_array_from_json(int32(), "[0, 0, 2, 4]");
        ASSERT_OK_AND_ASSIGN(auto list_arr, ListArray::from_arrays(*offsets_arr, *ext_arr));
        batch = RecordBatch::create(schema({field("f0", list(uuid()))}), 3, {list_arr});
        RoundtripBatch(batch, &read_batch);
        CompareBatch(*batch, *read_batch, false /* compare_metadata */);
    }

    TEST_F(TestExtensionType, UnrecognizedExtension) {
        auto ext_arr = ExampleUuid();
        auto batch = RecordBatch::create(schema({field("f0", uuid())}), 4, {ext_arr});

        auto storage_arr = static_cast<const ExtensionArray &>(*ext_arr).storage();

        // write full IPC stream including schema, then unregister type, then read
        // and ensure that a plain instance of the storage type is created
        ASSERT_OK_AND_ASSIGN(auto out_stream, io::BufferOutputStream::create());
        ASSERT_OK(ipc::write_record_batch_stream({batch}, ipc::IpcWriteOptions::defaults(),
            out_stream.get()));

        ASSERT_OK_AND_ASSIGN(auto complete_ipc_stream, out_stream->finish());

        ASSERT_OK(UnregisterExtensionType("uuid"));
        auto ext_metadata =
                key_value_metadata({
                    {"ARROW:extension:name", "uuid"},
                    {"ARROW:extension:metadata", "uuid-serialized"}
                });
        auto ext_field = field("f0", fixed_size_binary(16), true, ext_metadata);
        auto batch_no_ext = RecordBatch::create(schema({ext_field}), 4, {storage_arr});

        io::BufferReader reader(complete_ipc_stream);
        std::shared_ptr<RecordBatchReader> batch_reader;
        ASSERT_OK_AND_ASSIGN(batch_reader, ipc::RecordBatchStreamReader::open(&reader));
        std::shared_ptr<RecordBatch> read_batch;
        ASSERT_OK(batch_reader->read_next(&read_batch));
        CompareBatch(*batch_no_ext, *read_batch);
    }

    std::shared_ptr<Array> ExampleParametric(std::shared_ptr<DataType> type,
                                             const std::string &json_data) {
        auto arr = assert_array_from_json(int32(), json_data);
        auto ext_data = arr->data()->copy();
        ext_data->type = type;
        return make_array(ext_data);
    }

    TEST_F(TestExtensionType, ParametricTypes) {
        auto p1_type = std::make_shared<Parametric1Type>(6);
        auto p1 = ExampleParametric(p1_type, "[null, 1, 2, 3]");

        auto p2_type = std::make_shared<Parametric1Type>(12);
        auto p2 = ExampleParametric(p2_type, "[2, null, 3, 4]");

        auto p3_type = std::make_shared<Parametric2Type>(2);
        auto p3 = ExampleParametric(p3_type, "[5, 6, 7, 8]");

        auto p4_type = std::make_shared<Parametric2Type>(3);
        auto p4 = ExampleParametric(p4_type, "[5, 6, 7, 9]");

        ASSERT_OK(RegisterExtensionType(std::make_shared<Parametric1Type>(-1)));
        ASSERT_OK(RegisterExtensionType(p3_type));
        ASSERT_OK(RegisterExtensionType(p4_type));

        auto batch = RecordBatch::create(schema({
                                             field("f0", p1_type), field("f1", p2_type),
                                             field("f2", p3_type), field("f3", p4_type)
                                         }),
                                         4, {p1, p2, p3, p4});

        std::shared_ptr<RecordBatch> read_batch;
        RoundtripBatch(batch, &read_batch);
        CompareBatch(*batch, *read_batch, false /* compare_metadata */);
    }

    TEST_F(TestExtensionType, ParametricEquals) {
        auto p1_type = std::make_shared<Parametric1Type>(6);
        auto p2_type = std::make_shared<Parametric1Type>(6);
        auto p3_type = std::make_shared<Parametric1Type>(3);

        ASSERT_TRUE(p1_type->equals(p2_type));
        ASSERT_FALSE(p1_type->equals(p3_type));

        ASSERT_EQ(p1_type->fingerprint(), "");
    }

    std::shared_ptr<Array> ExampleStruct() {
        auto ext_type = std::make_shared<ExtStructType>();
        auto storage_type = ext_type->storage_type();
        auto arr = assert_array_from_json(storage_type, "[[1, 0.1], [2, 0.2]]");

        auto ext_data = arr->data()->copy();
        ext_data->type = ext_type;
        return make_array(ext_data);
    }

    TEST_F(TestExtensionType, ValidateExtensionArray) {
        auto ext_arr1 = ExampleUuid();
        auto p1_type = std::make_shared<Parametric1Type>(6);
        auto ext_arr2 = ExampleParametric(p1_type, "[null, 1, 2, 3]");
        auto ext_arr3 = ExampleStruct();
        auto ext_arr4 = ExampleComplex128();

        ASSERT_OK(ext_arr1->validate_full());
        ASSERT_OK(ext_arr2->validate_full());
        ASSERT_OK(ext_arr3->validate_full());
        ASSERT_OK(ext_arr4->validate_full());
    }
} // namespace nebula
